#include "spdlimit.h"
#include "ntdtime.h"

#define MAX_IO_EVER 16384

TokenBucketCfg::TokenBucketCfg(uint32 rate_limit, uint32 rate_burst, uint32 ms_per_tick)
{
	Init(rate_limit, rate_burst, ms_per_tick);
}

void TokenBucketCfg::Init(uint32 rate_limit, uint32 rate_burst, uint32 ms_per_tick)
{
	//convert Bytes/sec to Bytes per tick
	rate_limit_ = static_cast<uint64>(rate_limit) * ms_per_tick / 1000;
	rate_burst_ = static_cast<uint64>(rate_burst) * ms_per_tick / 1000;
	if( rate_burst_ < MSS_SIZE ){
		rate_burst_ = MSS_SIZE;
	}
	ms_per_tick_ = ms_per_tick;
}

TokenBucket::TokenBucket()
{
	cfg_ = NULL;
	tokens_left_ = 0;
	last_update_ = 0;
	lower_mark_ = 0;
	refill_timer_.AddListener(this);
}

TokenBucket::~TokenBucket()
{
	refill_timer_.RemoveListener(this);
	refill_timer_.Stop();
}

void TokenBucket::Init(const TokenBucketCfg* cfg, bool reset, int32 lower_mark)
{
	lower_mark_ = lower_mark;

	if( cfg != NULL ){
		if( cfg_ == NULL ){
			cfg_ = new TokenBucketCfg();
		}
		*cfg_ = *cfg;
	}else{
		delete cfg_;
		cfg_ = NULL;
		tokens_left_ = 0;
		last_update_ = 0;
		refill_timer_.Stop();
	}

	if( cfg_ ){
		if( !reset ){
			tokens_left_ = cfg_->rate_limit_;
			last_update_ = GetTicks();
		}else if( tokens_left_ > cfg_->rate_burst_ ){
			tokens_left_ = cfg_->rate_burst_;
		}
	}


}

int32 TokenBucket::GetLimit()
{
	if( cfg_ == NULL ){
		return MAX_IO_EVER;
	}

	Update();

	uint32 max_so_far = tokens_left_>0?tokens_left_:0;
	if( max_so_far > MAX_IO_EVER ){
		max_so_far = MAX_IO_EVER;
	}

	//return tokens_left_>0?tokens_left_:0;
	return max_so_far;
}

//return true if tokens are left, false if no tokens
//if no tokens, need to suspend io events notification until next refilled
//should start a refill timer at the same time
int32 TokenBucket::RemoveTokens(int32 bytes)
{
	if( cfg_ == NULL ){
		return MAX_IO_EVER;
	}

	tokens_left_ -= bytes;
	if( tokens_left_ <= lower_mark_ ){
		refill_timer_.Start(cfg_->ms_per_tick_, false);
		//notify to suspend io events
		std::for_each(listener_set_.begin(), listener_set_.end(), 
			std::bind2nd(std::mem_fun1(&ITokenBucketListner::OnTokenEmpty), this));
	}

	return tokens_left_;
}

//Refill tokens
void TokenBucket::OnTimer(NTDTimer* pTimer)
{
	Update();

	if( tokens_left_ > lower_mark_ ){
		//notify there're tokens now
		std::for_each(listener_set_.begin(), listener_set_.end(), 
			std::bind2nd(std::mem_fun1(&ITokenBucketListner::OnTokenFilled), this));
	}else{
		//refill in timer
		refill_timer_.Start(cfg_->ms_per_tick_, false);
	}
}

uint32 TokenBucket::GetTicks()
{
	return Now_ms() / cfg_->ms_per_tick_;
}
	
bool TokenBucket::Update()
{
	assert(cfg_);

	uint32 current_tick = GetTicks();

	uint32 n_ticks = current_tick - last_update_;

	/* Make sure some ticks actually happened, and that time didn't
	* roll back. */
	if (n_ticks == 0 || n_ticks > INT_MAX)
		return false;

	/* Naively, we would say
	bucket->limit += n_ticks * cfg->rate;

	if (bucket->limit > cfg->maximum)
	bucket->limit = cfg->maximum;

	But we're worried about overflow, so we do it like this:
	*/

	//LOG4CPLUS_TRACE(g_logger,"Refill tokens, n_ticks: "<<n_ticks);

	if ((cfg_->rate_burst_ - tokens_left_) / n_ticks < cfg_->rate_limit_){
		tokens_left_ = cfg_->rate_burst_;
	}else{
		tokens_left_ += n_ticks * cfg_->rate_limit_;
	}

	last_update_ = current_tick;

	return true;
}

void TokenBucket::AddListener(ITokenBucketListner * listener)
{
	listener_set_.insert(listener);
}

void TokenBucket::RemoveListener(ITokenBucketListner * listener)
{
	listener_set_.erase(listener);
}

//////////////////////////////////////////////////////////////////////////

//SpdLimit::SpdLimit()
//{
//	m_burstSize = 1500;
//	m_spdLimit = UINT_MAX;
//	m_msCircle = m_burstSize*1000/m_spdLimit;
//	m_leftBytes = 0;
//	m_lastPeriod = 0;
//}
//
//uint32 SpdLimit::GetLeftBytes(uint32 exceptedBytes)
//{
//	unsigned long period = GetTickCount()/m_msCircle;
//	if ( m_lastPeriod != period )
//	{
//		m_lastPeriod = period;
//		m_leftBytes = m_spdLimit*m_msCircle/1000;
//	}
//
//	uint32 result = 0;
//	result = min(exceptedBytes, m_leftBytes);
//	m_leftBytes -= result;
//
//	return result;
//}